% Generate Figure 3.
% It is advised to run the penalized Bierens test separately from the other tests, as it may take considerable amount of time.
% Figure 3 in Chen et al. (2025) was generated using data simulated on the HPC at Columbia University, which uses distributed seed numbers.
% There may be simulation errors in the final results if different seed numbers are used.

clear
clc

% Designate working directories
dir = '.../Replication package/';
cd(dir);
addpath(genpath(strcat(dir,"Code/src")))
addpath(genpath(strcat(dir,"Datasets")))

% Set seed
seed = 4534546 ;
rng(seed,'twister');

alpha_level = 0.1;          % nominal level of the tests
theta_grid = -0.6:0.1:0.6;
bnd = 5 ;                   % bnd > 0 defines the search space as [-bndbnd]^p, where p = dim(gamma).
bR  = 5000 ;                % bootstrap replication number used in grid search method.
sz_Gamma = 200000 ;         % number of grids in grid search method.
splits = 1000 ;             % reduce memory requirement
lambda_grid = 1:(-0.1):0 ;  % grid for optimal lambda search.
B = 2 ;

% Load the data
data = double(load("USAY.mat").dataset);
lag = 2;  % use two-period-lagged variables as instruments
dat = rmmissing([data(:, 6), data(:, 8), lagmatrix(data(:, 3:5), 2), lagmatrix(data(:, 6), 2)]);
Y = dat(:, 1);
X = dat(:, 2);
Z = dat(:, 3:6);

[pi, ~, resid_fs] = regress(X, [ones(105,1) Z]);
[pi_res, ~, vhat] = regress(resid_fs, [ones(105,1) Z Z.^2]);
Xhat = [ones(105,1) Z]*pi;
theta  = regress(Y, [ones(105,1) Xhat]);
uhat = Y - [ones(105,1) X] * theta;
sigma_u = std(uhat);
sigma_v = std(vhat);
rho = 0.8;

Z1 = lagmatrix(Z(:,1:3), 2);

simul_input.T = 105;
simul_input.lag = lag;
simul_input.pi = pi;
simul_input.pi_res = pi_res;
simul_input.rho = rho;
simul_input.theta = theta;
simul_input.sigma_u = sigma_u;
simul_input.sigma_v = sigma_v;
simul_input.X = X;
simul_input.Y = Y;
simul_input.Z1 = Z1;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Simulations
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
h = @(y, theta) y(:,1) - y(:,2)*theta;
hdot = @(y, theta) -y(:,2);
theta_grid = theta_grid + theta(2);  % translate by theta(2) = -0.0287
cv_AR = chi2inv(1-alpha_level, 4);

R = 50;  % replication number

rec_rej        = zeros(R, length(lambda_grid), length(theta_grid));
rec_bspv       = zeros(R, length(lambda_grid), length(theta_grid));
rec_opt_lambda = zeros(R, length(theta_grid));
rec_rtime_pB   = zeros(R, 1);
rec_opt_pv     = zeros(R, length(theta_grid));
rec_opt_rej    = zeros(R, length(theta_grid));

for rep = 1:R
    data = gen_data_NL(simul_input, 2);
    y = data.y; dy = y - mean(y);
    x = data.x; dx = x - mean(x);
    z = data.z; dz = z - mean(z);

    u = y - mean(y) - (x - mean(x))*theta_grid;
    test = pen_Bierens_test_GS(u, z, bnd, lambda_grid, "sz_Gamma", sz_Gamma, "splits", splits, ...
                                    "bR", bR, 'gradient', -x+mean(x), 'B', B, 'dm', 1);

    rec_rej(rep, :, :)     = test.rejection;
    rec_bspv(rep, :, :)    = test.bs_pv;
    rec_opt_lambda(rep, :) = test.opt_lambda;
    rec_opt_pv(rep, :)     = test.opt_pv;
    rec_opt_rej(rep, :)    = (test.opt_pv <= alpha_level);
    rec_rtime_pB(rep)      = test.runtime;

    % Dominguez-Lobato
    test_DL = Dominguez_Lobato([dy, dx], z, h, hdot , 1, 'lb', -20, 'ub', 20);
    rec_rej_DL(rep, :) = (theta_grid < test_DL.lowerCI) | (theta_grid > test_DL.upperCI);
    rec_rtime_DL(rep) = test_DL.rtime;

    % AR test (NOT heteroscedasticity robust)
    Tstart = tic();
    P_Z = dz / (dz' * dz) * dz';
    M_Z = eye(size(z,1)) - P_Z;
    AR_stat = zeros(length(theta_grid), 1);
    for i = 1:length(theta_grid)
        AR_stat(i) = (dy - dx * theta_grid(i))' * P_Z * (dy - dx * theta_grid(i)) ./ ((dy - dx * theta_grid(i))' * M_Z * (dy - dx * theta_grid(i)) / (size(z,1) - size(z,2)));
    end
    rec_rej_AR(rep, :) = (AR_stat > cv_AR);
    Tend = toc(Tstart);
    rec_rtime_AR(rep) = Tend;

    % 2SLS
    Tstart = tic();
    theta_2sls = (dx' * P_Z * dx) \ (dx' * P_Z * dy);
    ehat = dy - dx * theta_2sls;
    Vhat = (dx' * P_Z * dx) \ (dx' * P_Z * diag(ehat.^2) * P_Z * dx) / (dx' * P_Z * dx);  % Heteroscedasticity-consistent variance estimator
    se = sqrt(Vhat);
    lowerCI = theta_2sls - se * norminv(1-alpha_level/2);
    upperCI = theta_2sls + se * norminv(1-alpha_level/2);
    rec_rej_2SLS(rep, :) = (theta_grid < lowerCI) | (theta_grid > upperCI);
    Tend = toc(Tstart);
    rec_rtime_2SLS(rep) = Tend;
end

unpenalized_pv = squeeze(rec_bspv(:,end,:));
opt_pb = mean(rec_opt_rej);
unp_pb  = mean(unpenalized_pv <= .1);
unp_pb_adj = mean(unpenalized_pv <= .1) - unp_pb(7) + .1; unp_pb_adj = max(min(unp_pb_adj,1),0);
DL = mean(rec_rej_DL);
AR = mean(rec_rej_AR);
TSLS = mean(rec_rej_2SLS);

theta_grid = -0.6:0.1:0.6;

fig = figure('Visible', 'off');
hold on
plot(theta_grid, opt_pb, '-o', 'LineWidth', 1.5, 'Color', "#0072BD")
plot(theta_grid, unp_pb, '-x', 'LineWidth', 1.5, 'Color', "#D95319")
plot(theta_grid, unp_pb_adj, '-+', 'LineWidth', 1.0, 'Color', "#EDB120")
plot(theta_grid, AR, '-s', 'LineWidth', 1.5)
plot(theta_grid, TSLS, '-d', 'LineWidth', 1.5)
plot(theta_grid, DL, '-*', 'LineWidth', 1.5)
plot(theta_grid, 0.1*ones(length(theta_grid), 1),  '--', 'LineWidth', 0.5, 'Color', [.7 .7 .7])
ylabel('Probability of Rejection')
ylim([0.00, 1])
xlabel('$\theta_1 - \theta_{10}$', 'Interpreter', 'latex')
lgd = legend("Opt.", "Unp.", "Unp. w/ shift", "AR", "2SLS", "DL",'Interpreter', 'latex', 'Location', 'best')
lgd.Box = 'off';
saveas(fig, strcat(dir,"Figure_3.jpg"))

% Save the data for Table 4.
save(strcat(dir,'selected_lambdas.mat'), 'rec_opt_lambda')